"""

@Time: 2021/2/19 1:53
@Author:
@File: ml_controller.py
"""
from fastapi import APIRouter, Depends, Path
from motor.motor_asyncio import AsyncIOMotorDatabase
from bson.objectid import ObjectId

from models.dto import RESTfulModel
from models.dto.ml_dto import MLModelDetailDto
from dependencies import verify_token, get_mongo_db
from config import mongo_conf

# TODO: 改成文件管理模型
ml_router = APIRouter(tags=['ml'])


@ml_router.get('/model_detail/{model_id}',
               response_model=RESTfulModel[MLModelDetailDto],
#               dependencies=[Depends(verify_token), ],
               summary='根据一个训练模型的id来获取该模型的详细信息',
               description='该 id 是存于 MongoDB 中的 `_id` 值')
async def get_one_model_detail(model_id: str = Path(..., title='训练模型的 ID 值'),
                               mongo_db: AsyncIOMotorDatabase = Depends(get_mongo_db)):
    mongo_collection = mongo_db[mongo_conf.ML_MODEL_TEMP_COLLECTION_NAME]
    model: dict = await mongo_collection.find_one({'_id': ObjectId(model_id)})  # 根据文档 ID 找出该文档
    model['id'] = str(model['_id'])  # 将 ObjectID 转化成 str
    model.pop('_id')
    return RESTfulModel(code=0, data=model)


@ml_router.get('/save_model/{model_id}',
               response_model=RESTfulModel[str],
               dependencies=[Depends(verify_token), ],
               summary='将存储于暂存区的模型保存到保存区',
               description='该 id 是存于 MongoDB 中的 `_id` 值')
async def save_one_model(model_id: str = Path(..., title='训练模型的 ID 值'),
                         mongo_db: AsyncIOMotorDatabase = Depends(get_mongo_db)):
    mongo_temp_collection = mongo_db[mongo_conf.ML_MODEL_TEMP_COLLECTION_NAME]  # mongo 暂存区的集合
    mongo_saved_collection = mongo_db[mongo_conf.ML_MODEL_SAVED_COLLECTION_NAME]  # mongo 保存区的集合
    model_doc = await mongo_temp_collection.find_one_and_delete({'_id': ObjectId(model_id)})  # TODO: 是否应该delete？还是只需要find_one即可？
    await mongo_saved_collection.insert_one(model_doc)
    return RESTfulModel(code=0, data='success')




